#include <bits/stdc++.h>
using namespace std;

/*
3354. 使数组元素等于零
已解答
简单
相关标签
premium lock icon
相关企业
提示
给你一个整数数组 nums 。

开始时，选择一个满足 nums[curr] == 0 的起始位置 curr ，并选择一个移动 方向 ：向左或者向右。

此后，你需要重复下面的过程：

如果 curr 超过范围 [0, n - 1] ，过程结束。
如果 nums[curr] == 0 ，沿当前方向继续移动：如果向右移，则 递增 curr ；如果向左移，则 递减 curr 。
如果 nums[curr] > 0:
将 nums[curr] 减 1 。
反转 移动方向（向左变向右，反之亦然）。
沿新方向移动一步。
如果在结束整个过程后，nums 中的所有元素都变为 0 ，则认为选出的初始位置和移动方向 有效 。

返回可能的有效选择方案数目。

 

示例 1：

输入：nums = [1,0,2,0,3]

输出：2

解释：

可能的有效选择方案如下：

选择 curr = 3 并向左移动。
[1,0,2,0,3] -> [1,0,2,0,3] -> [1,0,1,0,3] -> [1,0,1,0,3] -> [1,0,1,0,2] -> [1,0,1,0,2] -> [1,0,0,0,2] -> [1,0,0,0,2] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,0].
选择 curr = 3 并向右移动。
[1,0,2,0,3] -> [1,0,2,0,3] -> [1,0,2,0,2] -> [1,0,2,0,2] -> [1,0,1,0,2] -> [1,0,1,0,2] -> [1,0,1,0,1] -> [1,0,1,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [0,0,0,0,0].
示例 2：

输入：nums = [2,3,4,0,4,1,0]

输出：0

解释：

不存在有效的选择方案。

 
*/

// 法一
class Solution {
public:
    // helper
    /// 从 cur start dir--> check all to 0(bool)
    bool isValid(vector<int> nums, int cur, int dir) {
        int n = nums.size();
        while (true) {
            if (cur < 0 || cur >= n)        break;  // 越界 
            if (nums[cur] == 0) cur += dir;
            else {
                nums[cur]--;
                dir *= -1;      // 反方向
                cur += dir;
            }
        }
        // check all
        for (int num : nums) {
            if (num != 0)   return false;
        }
        return true;
    }

    int countValidSelections(vector<int>& nums) {
        int n = nums.size();
        vector<int> zeros;      // 全部为0的集合
        for (int i = 0; i < n; i++) {
            if (nums[i] == 0)   zeros.push_back(i);
        }
        // init no pos --> return 0
        if (zeros.empty())  return 0;

        int ans = 0;
        // try 2 dirs
        for (int cur : zeros) {
            if (isValid(nums, cur, 1))  ans++;
            if (isValid(nums, cur, -1))  ans++;
        }
        return ans;
    }
};

// 法二
class Solution {
public:
    int countValidSelections(vector<int>& nums) {
        int total = reduce(nums.begin(), nums.end()); // 所有元素总和（总需减的次数）
        int ans = 0, pre = 0; // pre：累计到当前位置的元素和（左侧和）
        
        for (int x : nums) {
            if (x != 0)   pre += x; // 累加非0元素，计算左侧和
            else { // 当前位置是0（可能的初始位置）
                // 右侧和 = total - pre
                if (pre * 2 == total)                   ans += 2;        // 左侧和 = 右侧和，左右方向均有效                    
                else if (abs(pre * 2 - total) == 1)     ans++;              // 左右和差1，仅一个方向有效     
            }
        }
        return ans;
    }
};
